package org.wutopia.labcloud.domain.gateway.filter;

import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.UUID;

@Component
@Slf4j
public class LoggingFilter implements GlobalFilter, Ordered {

    private static final String REQUEST_START_TIME = "requestStartTime";
    private static final String REQUEST_ID = "requestId";

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 生成请求ID并记录请求开始时间
        String requestId = UUID.randomUUID().toString().replace("-", "");
        exchange.getAttributes().put(REQUEST_ID, requestId);
        exchange.getAttributes().put(REQUEST_START_TIME, System.currentTimeMillis());

        ServerHttpRequest request = exchange.getRequest();
        
        // 只有POST、PUT等有请求体的方法才需要记录请求体
        if (HttpMethod.POST.equals(request.getMethod()) || 
            HttpMethod.PUT.equals(request.getMethod()) || 
            HttpMethod.PATCH.equals(request.getMethod())) {
            return logRequestBody(exchange, chain, requestId);
        } else {
            // 记录请求进入信息
            logRequestIncoming(exchange, requestId, null);
            
            // 记录响应信息
            ServerHttpResponse response = exchange.getResponse();
            DataBufferFactory bufferFactory = response.bufferFactory();
            ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(response) {
                @Override
                public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                    if (body instanceof Flux) {
                        Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                        return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                            DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                            DataBuffer join = dataBufferFactory.join(dataBuffers);
                            byte[] content = new byte[join.readableByteCount()];
                            join.read(content);
                            DataBufferUtils.release(join);
                            String responseBody = new String(content, StandardCharsets.UTF_8);
                            logResponseOutgoing(exchange, requestId, responseBody);
                            return bufferFactory.wrap(content);
                        }));
                    }
                    return super.writeWith(body);
                }
            };
            
            return chain.filter(exchange.mutate().response(decoratedResponse).build())
                    .then(Mono.fromRunnable(() -> logRequestOutgoing(exchange, requestId)));
        }
    }

    private Mono<Void> logRequestBody(ServerWebExchange exchange, GatewayFilterChain chain, String requestId) {
        ServerHttpRequest serverHttpRequest = exchange.getRequest();
        
        return DataBufferUtils.join(serverHttpRequest.getBody())
                .flatMap(dataBuffer -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    
                    String requestBody = new String(bytes, StandardCharsets.UTF_8);
                    
                    // 记录请求进入信息（包含请求体）
                    logRequestIncoming(exchange, requestId, requestBody);
                    
                    // 重新包装请求体，确保后续过滤器可以读取
                    ServerHttpRequestDecorator decoratedRequest = new ServerHttpRequestDecorator(serverHttpRequest) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return Flux.just(new DefaultDataBufferFactory().wrap(bytes));
                        }
                    };
                    
                    // 包装响应以记录响应体
                    ServerHttpResponse response = exchange.getResponse();
                    DataBufferFactory bufferFactory = response.bufferFactory();
                    ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(response) {
                        @Override
                        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                            if (body instanceof Flux) {
                                Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                                return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                                    DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                                    DataBuffer join = dataBufferFactory.join(dataBuffers);
                                    byte[] content = new byte[join.readableByteCount()];
                                    join.read(content);
                                    DataBufferUtils.release(join);
                                    String responseBody = new String(content, StandardCharsets.UTF_8);
                                    logResponseOutgoing(exchange, requestId, responseBody);
                                    return bufferFactory.wrap(content);
                                }));
                            }
                            return super.writeWith(body);
                        }
                    };
                    
                    return chain.filter(exchange.mutate().request(decoratedRequest).response(decoratedResponse).build())
                            .then(Mono.fromRunnable(() -> logRequestOutgoing(exchange, requestId)));
                });
    }

    private void logRequestIncoming(ServerWebExchange exchange, String requestId, String requestBody) {
        String method = exchange.getRequest().getMethod().name();
        String path = exchange.getRequest().getURI().getPath();
        String query = exchange.getRequest().getURI().getQuery();
        String remoteAddress = exchange.getRequest().getRemoteAddress() != null ?
                exchange.getRequest().getRemoteAddress().toString() : "Unknown";

        log.info(">>> [{}] {} {}?{} from {}", requestId, method, path, query != null ? query : "", remoteAddress);
        
        if (requestBody != null && !requestBody.isEmpty()) {
            log.info(">>> [{}] Request Body: {}", requestId, requestBody);
        }
    }

    private void logRequestOutgoing(ServerWebExchange exchange, String requestId) {
        long startTime = exchange.getAttributeOrDefault(REQUEST_START_TIME, System.currentTimeMillis());
        long executeTime = System.currentTimeMillis() - startTime;
        String path = exchange.getRequest().getURI().getPath();
        int statusCode = exchange.getResponse().getStatusCode() != null ?
                exchange.getResponse().getStatusCode().value() : 0;

        log.info("<<< [{}] {} {} {}ms", requestId, statusCode, path, executeTime);
    }

    private void logResponseOutgoing(ServerWebExchange exchange, String requestId, String responseBody) {
        int statusCode = exchange.getResponse().getStatusCode() != null ?
                exchange.getResponse().getStatusCode().value() : 0;
        log.info("<<< [{}] {} Response Body: {}", requestId, statusCode, responseBody);
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 1; // 确保在其他过滤器之前执行
    }
}
